"""
@author: chenzhenhua
@project: xltpl
@file: office_exam.py
@time: 2021/6/16 0016 13:48
@desc:
"""

import os

import numpy as np
import pandas as pd
from xltpl.writerx import BookWriter as BookWriterx


def get_all_path(dir):
    """
    获取文件夹下所有文件名

    :param dir:文件路径
    :return:
    """
    return os.listdir(dir)


class ArrangeExam(object):
    def __init__(self, dir, subject, ever_num, path_index, addr_code):
        self.dir = dir
        self.subject = subject
        self.ever_num = ever_num
        self.path_index = path_index
        self.addr_code = addr_code

    def base_data(self):
        """
        读取所有数据

        :return:所有dataframe数据合并到一个pandas
        """

        data = []
        path = get_all_path(self.dir)
        path = np.array(path)[self.path_index]
        for pathi in path:
            path_xlsx = os.path.join(self.dir, pathi)
            datai = pd.read_excel(path_xlsx)[['姓名', '班级', '学号']]
            data.append(datai)
        data_all = pd.concat(data)
        data_all['学号'] = data_all['学号'].astype(str)
        return data_all

    def split_data(self, n1, n2, data_arr):
        """
        切分数据

        :param n1:每个考场从n1位置开始切分
        :param n2:每个考场从n2位置切分截止
        :param data_arr:base_data返回数据的numpy格式
        :return:表头是'座位号', '课程名称', '班级', '学号', '姓名'的dataframe数据
        """
        data = pd.DataFrame(data_arr[n1:n2, :], columns=['姓名', '班级', '学号'])
        data['学号'] = data['学号'].astype(str)

        list_index = list(range(n2 - n1))
        subject_name = [self.subject] * (n2 - n1)

        data['座位号'] = list_index
        data['座位号'] = data['座位号'].astype(str)
        data['课程名称'] = subject_name

        data = data[['座位号', '课程名称', '班级', '学号', '姓名']]
        return data

    def add_useful_info(self):
        """

        :return:返回每个考场信息，保存在list中
        """
        data_all = self.base_data()
        data_arr = data_all.values
        data_final = []
        s = 0
        for i in range(len(self.ever_num)):
            n1 = s
            s = s + self.ever_num[i]
            n2 = s
            datai = self.split_data(n1, n2, data_arr)
            data_final.append(datai)
        return data_final

    def write_to_excel(self, path_origin, path_new):
        """
        写入excel文件的单个sheet中

        :param path_origin:模板文件路径
        :param path_new:保存文件路径
        :return:
        """
        writer = BookWriterx(path_origin)
        data_to_excel = self.add_useful_info()

        rows_all = []
        for i in range(len(data_to_excel)):
            n = len(data_to_excel[i])
            set_nums = list(range(1, n + 1))
            subjiect_names = [self.subject] * n
            classes = data_to_excel[i]['班级'].values
            student_nums = data_to_excel[i]['学号'].values
            names = data_to_excel[i]['姓名'].values
            sigs = [''] * n
            scores = [''] * n
            f_scores = [''] * n
            rows = np.array([set_nums, subjiect_names, classes, student_nums, names, sigs, scores, f_scores]).T
            rows_all.append([rows, self.addr_code[i]])

        exam = {}
        exam['rows_all'] = rows_all

        exam['sheet_name'] = self.subject
        payloads = [exam]
        writer.render_book(payloads=payloads)

        writer.save(path_new)

    def write_to_excels(self, path_origin, path_new):
        """
        写入excel的多个sheet中

        :param path_origin: 模板文件路径
        :param path_new: 写入路径
        :return:
        """
        writer = BookWriterx(path_origin)
        data_to_excel = self.add_useful_info()

        payloads = []
        for i in range(len(data_to_excel)):
            n = len(data_to_excel[i])
            set_nums = list(range(1, n + 1))
            subjiect_names = [self.subject] * n
            classes = data_to_excel[i]['班级'].values
            student_nums = data_to_excel[i]['学号'].values
            names = data_to_excel[i]['姓名'].values
            sigs = [''] * n
            scores = [''] * n
            f_scores = [''] * n
            rows = np.array([set_nums, subjiect_names, classes, student_nums, names, sigs, scores, f_scores]).T
            exam = {}
            exam['rows'] = rows
            exam['addr_code'] = self.addr_code[i]
            exam['sheet_name'] = str(i + 1)
            payloads.append(exam)
        writer.render_book(payloads=payloads)

        writer.save(path_new)


class FileClass(ArrangeExam):
    def base_data(self):
        """
        读取所有数据,班级名和文件的班级名一致

        :return:所有dataframe数据合并到一个pandas
        """

        data = []
        path = get_all_path(self.dir)
        path = np.array(path)[self.path_index]
        for pathi in path:
            name = pathi.replace('.xlsx', '').replace('.xls', '')
            if '20本科' not in name:
                name = '20本科' + name
            path_xlsx = os.path.join(self.dir, pathi)
            datai = pd.read_excel(path_xlsx)[['姓名', '学号']]
            n = len(datai)
            this_class = [name] * n
            datai['班级'] = np.array(this_class)
            data.append(datai)
        data_all = pd.concat(data)
        data_all['学号'] = data_all['学号'].astype(str)
        return data_all
